import pygame
import random
import heapq
from collections import deque
import sys

class Snake:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.reset()
    
    def reset(self):
        self.body = [(self.width // 2, self.height // 2)]
        self.direction = (1, 0)
        self.grow = False
    
    def move(self, new_direction=None):
        if new_direction:
            self.direction = new_direction
        
        head = self.body[0]
        new_head = (head[0] + self.direction[0], head[1] + self.direction[1])
        
        if self.grow:
            self.body.insert(0, new_head)
            self.grow = False
        else:
            self.body.insert(0, new_head)
            self.body.pop()
    
    def check_collision(self):
        head = self.body[0]
        
        if head[0] < 0 or head[0] >= self.width or head[1] < 0 or head[1] >= self.height:
            return True
        
        if head in self.body[1:]:
            return True
        
        return False
    
    def get_head(self):
        return self.body[0]
    
    def get_body(self):
        return self.body

class Food:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.position = None
        self.spawn()
    
    def spawn(self, snake_body=None):
        if snake_body is None:
            snake_body = []
        
        available_positions = []
        for x in range(self.width):
            for y in range(self.height):
                if (x, y) not in snake_body:
                    available_positions.append((x, y))
        
        if available_positions:
            self.position = random.choice(available_positions)
    
    def get_position(self):
        return self.position

class AStar:
    def __init__(self, width, height):
        self.width = width
        self.height = height
    
    def heuristic(self, a, b):
        return abs(a[0] - b[0]) + abs(a[1] - b[1])
    
    def get_neighbors(self, pos, snake_body):
        x, y = pos
        neighbors = []
        
        for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
            new_x, new_y = x + dx, y + dy
            
            if (0 <= new_x < self.width and 
                0 <= new_y < self.height and 
                (new_x, new_y) not in snake_body[1:]):
                neighbors.append((new_x, new_y))
        
        return neighbors
    
    def find_path(self, start, goal, snake_body):
        if not goal:
            return None
        
        open_set = []
        heapq.heappush(open_set, (0, start))
        
        came_from = {}
        g_score = {start: 0}
        f_score = {start: self.heuristic(start, goal)}
        
        while open_set:
            current = heapq.heappop(open_set)[1]
            
            if current == goal:
                path = []
                while current in came_from:
                    path.append(current)
                    current = came_from[current]
                return path[::-1]
            
            for neighbor in self.get_neighbors(current, snake_body):
                tentative_g_score = g_score[current] + 1
                
                if neighbor not in g_score or tentative_g_score < g_score[neighbor]:
                    came_from[neighbor] = current
                    g_score[neighbor] = tentative_g_score
                    f_score[neighbor] = tentative_g_score + self.heuristic(neighbor, goal)
                    heapq.heappush(open_set, (f_score[neighbor], neighbor))
        
        return None

class SnakeGame:
    def __init__(self, width=20, height=15, cell_size=30):
        pygame.init()
        
        self.width = width
        self.height = height
        self.cell_size = cell_size
        
        self.screen_width = width * cell_size
        self.screen_height = height * cell_size
        
        self.screen = pygame.display.set_mode((self.screen_width, self.screen_height))
        pygame.display.set_caption("贪吃蛇 - A*算法自动寻路")
        
        self.clock = pygame.time.Clock()
        self.font = pygame.font.Font(None, 36)
        
        self.snake = Snake(width, height)
        self.food = Food(width, height)
        self.astar = AStar(width, height)
        
        self.score = 0
        self.game_over = False
        self.auto_mode = True
        self.path = []
        self.path_update_counter = 0
        
        self.colors = {
            'black': (0, 0, 0),
            'white': (255, 255, 255),
            'green': (0, 255, 0),
            'red': (255, 0, 0),
            'blue': (0, 0, 255),
            'gray': (128, 128, 128)
        }
    
    def get_direction_from_path(self):
        if not self.path or len(self.path) < 2:
            return None
        
        head = self.snake.get_head()
        next_pos = self.path[0]
        
        dx = next_pos[0] - head[0]
        dy = next_pos[1] - head[1]
        
        if dx == 1:
            return (1, 0)
        elif dx == -1:
            return (-1, 0)
        elif dy == 1:
            return (0, 1)
        elif dy == -1:
            return (0, -1)
        
        return None
    
    def update_path(self):
        if not self.auto_mode:
            return
        
        head = self.snake.get_head()
        food_pos = self.food.get_position()
        
        self.path = self.astar.find_path(head, food_pos, self.snake.get_body())
        
        if not self.path:
            safe_moves = self.astar.get_neighbors(head, self.snake.get_body())
            if safe_moves:
                self.path = [safe_moves[0]]
    
    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    self.auto_mode = not self.auto_mode
                elif event.key == pygame.K_r and self.game_over:
                    self.reset_game()
                elif not self.auto_mode:
                    head = self.snake.get_head()
                    if event.key == pygame.K_UP and self.snake.direction != (0, 1):
                        self.snake.direction = (0, -1)
                    elif event.key == pygame.K_DOWN and self.snake.direction != (0, -1):
                        self.snake.direction = (0, 1)
                    elif event.key == pygame.K_LEFT and self.snake.direction != (1, 0):
                        self.snake.direction = (-1, 0)
                    elif event.key == pygame.K_RIGHT and self.snake.direction != (-1, 0):
                        self.snake.direction = (1, 0)
        
        return True
    
    def update(self):
        if self.game_over:
            return
        
        if self.auto_mode:
            self.path_update_counter += 1
            
            if self.path_update_counter >= 3:
                self.update_path()
                self.path_update_counter = 0
            
            direction = self.get_direction_from_path()
            if direction:
                self.snake.direction = direction
            
            if self.path and len(self.path) > 1:
                next_pos = self.path[0]
                if next_pos == self.snake.get_head():
                    self.path.pop(0)
        
        self.snake.move()
        
        if self.snake.check_collision():
            self.game_over = True
            return
        
        head = self.snake.get_head()
        if head == self.food.get_position():
            self.snake.grow = True
            self.food.spawn(self.snake.get_body())
            self.score += 10
            self.path = []
    
    def draw(self):
        self.screen.fill(self.colors['black'])
        
        for x in range(self.width + 1):
            pygame.draw.line(self.screen, self.colors['gray'], 
                           (x * self.cell_size, 0), 
                           (x * self.cell_size, self.screen_height))
        
        for y in range(self.height + 1):
            pygame.draw.line(self.screen, self.colors['gray'], 
                           (0, y * self.cell_size), 
                           (self.screen_width, y * self.cell_size))
        
        for i, segment in enumerate(self.snake.get_body()):
            x, y = segment
            rect = pygame.Rect(x * self.cell_size + 1, y * self.cell_size + 1, 
                             self.cell_size - 2, self.cell_size - 2)
            
            if i == 0:
                pygame.draw.rect(self.screen, self.colors['blue'], rect)
            else:
                pygame.draw.rect(self.screen, self.colors['green'], rect)
        
        food_pos = self.food.get_position()
        if food_pos:
            food_rect = pygame.Rect(food_pos[0] * self.cell_size + 1, 
                                  food_pos[1] * self.cell_size + 1,
                                  self.cell_size - 2, self.cell_size - 2)
            pygame.draw.rect(self.screen, self.colors['red'], food_rect)
        
        if self.auto_mode and self.path:
            for i, pos in enumerate(self.path[:5]):
                x, y = pos
                center = (x * self.cell_size + self.cell_size // 2,
                         y * self.cell_size + self.cell_size // 2)
                pygame.draw.circle(self.screen, self.colors['white'], center, 3)
        
        score_text = self.font.render(f"Score: {self.score}", True, self.colors['white'])
        self.screen.blit(score_text, (10, 10))
        
        mode_text = "AUTO" if self.auto_mode else "MANUAL"
        mode_surface = self.font.render(f"Mode: {mode_text}", True, self.colors['white'])
        self.screen.blit(mode_surface, (10, 50))
        
        if self.game_over:
            game_over_text = self.font.render("GAME OVER! Press R to restart", True, self.colors['red'])
            text_rect = game_over_text.get_rect(center=(self.screen_width // 2, self.screen_height // 2))
            self.screen.blit(game_over_text, text_rect)
        
        instructions = [
            "SPACE: Toggle Auto/Manual mode",
            "R: Restart (when game over)",
            "Arrow keys: Manual control"
        ]
        
        for i, instruction in enumerate(instructions):
            text = pygame.font.Font(None, 24).render(instruction, True, self.colors['white'])
            self.screen.blit(text, (self.screen_width - 300, 10 + i * 25))
        
        pygame.display.flip()
    
    def reset_game(self):
        self.snake.reset()
        self.food.spawn()
        self.score = 0
        self.game_over = False
        self.path = []
        self.path_update_counter = 0
    
    def run(self):
        running = True
        
        while running:
            running = self.handle_events()
            self.update()
            self.draw()
            self.clock.tick(10)
        
        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    game = SnakeGame()
    game.run()